import argparse
import time

from audio_text_common import gen_subtitle
from llm import call_llm
from llm import gen_subtitle_check_prompt
from log_manager import LogManager
from srt_formatter import format_srt


def do_task(
    video_path,
    output_path,
    ref_path,
    split_chars="，。,.；;",
    type="local",
    model_name="medium",
):
    """
    生成字幕
    Args:
        video_path: 视频文件路径
        output_path: 输出的 SRT 文件路径
        model_name: whisper 模型名称
    """
    if ref_path:
        with open(ref_path, "r", encoding="utf-8") as f:
            ref_text = f.read()
    else:
        ref_text = None
    # 首先使用 whisper 提取语音文本
    subpath = output_path.rsplit(".", 1)[0]
    whisper_output_path = subpath + "_whisper.srt"
    LogManager.log("开始生成原始字幕...")
    subtitle_text = gen_subtitle(
        video_path, whisper_output_path, split_chars, type, model_name
    )
    LogManager.log(f"原始字幕文件已保存到:{whisper_output_path}")

    # 读取生成的字幕文件
    LogManager.log("字幕生成完成，开始校对")
    system_prompt, user_prompt = gen_subtitle_check_prompt(subtitle_text, ref_text)
    # 校对字幕
    corrected_subtitle = call_llm(system_prompt, user_prompt)
    LogManager.log("校对完成，开始格式化")
    corrected_subtitle_path = subpath + "_corrected.srt"
    with open(corrected_subtitle_path, "w", encoding="utf-8") as f:
        f.write(corrected_subtitle)
    # 格式化字幕
    improved_subtitles = format_srt(corrected_subtitle, split_chars)
    LogManager.log("格式化完成，开始保存")
    # 保存优化后的字幕
    with open(output_path, "w", encoding="utf-8") as f:
        f.write(improved_subtitles)
    LogManager.log(f"优化后的字幕已保存到:{output_path}")


if __name__ == "__main__":
    start_time = time.time()
    parser = argparse.ArgumentParser(description="字幕识别及校验")
    parser.add_argument("video", help="输入视频文件路径")
    parser.add_argument("--ref", default=None, help="原始文本文件路径")
    parser.add_argument("--output", default=None, help="输出的SRT文件路径")
    parser.add_argument("--split", default="，。,.；;", help="字幕分隔字符")
    parser.add_argument("--model", default="small", help="whisper模型名称")
    parser.add_argument("--type", default="local", help="字幕生成方式")
    args = parser.parse_args()
    ref_path = args.ref
    video_path = args.video
    if args.split:
        split_chars = args.split
    else:
        split_chars = "，。,.；;"
    model = args.model
    if args.output:
        output = args.output
    else:
        output = args.video.rsplit(".", 1)[0] + "_subtitles.srt"
    do_task(video_path, output, ref_path, split_chars, args.type, model)
    end_time = time.time()
    total_seconds = end_time - start_time
    minutes = int(total_seconds // 60)
    seconds = total_seconds % 60
    LogManager.log(f"总运行时间: {minutes}分钟 {seconds:.2f}秒")
